import argparse
import os
import glob
import cv2
import numpy as np


def parse_args():
    parser = argparse.ArgumentParser(description='Using template matching for license plate character recognition')
    # add argument as you like
    parser.add_argument('--img-dir', default='/work/dn/week3/work_dirs/test_clipped_resize_char_edge',
                        help='License plate image directory to be recognized')
    parser.add_argument('--tmp-dir', default='/work/dn/week3/work_dirs/Train_resized_edge')
    parser.add_argument('--method', default='sqdiff', help='the match method')

    args = parser.parse_args()
    return args


method_dict = {
    'sqdiff': cv2.TM_SQDIFF_NORMED,
    'ccorr': cv2.TM_CCORR_NORMED,
    'ccoeff': cv2.TM_CCOEFF_NORMED
}

if __name__ == '__main__':
    args = parse_args()
    assert args.method in method_dict
    # ========用三种方法进行模板匹配========
    method = method_dict[args.method]
    print(f'parsing image from {args.img_dir}...use the method: {args.method}')

    # 首先读取十张模板。
    templates = []
    for i in range(10):
        tem_dir = os.path.join(args.tmp_dir, 'sobel-'+str(i)+'.jpg')
        tem_img = cv2.imread(tem_dir, cv2.IMREAD_GRAYSCALE)   # h, w
        templates.append(tem_img)
    h, w = templates[0].shape

    all = 0           # 总样本数
    right = 0         # 识别正确的数目

    for test_img_dir in glob.glob(os.path.join(args.img_dir, '*.jpg')):
        gt = int(test_img_dir.split('/')[-1].split('.')[0].split('-')[-1])
        test_img = cv2.resize(cv2.imread(test_img_dir, cv2.IMREAD_GRAYSCALE), (w, h), interpolation=cv2.INTER_NEAREST)

        res = []   # 记录对十个模板的匹配程度
        for tem in templates:
            result = cv2.matchTemplate(test_img, tem, method)
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
            if method == cv2.TM_SQDIFF_NORMED:
                res.append(min_val)
            else:
                res.append(max_val)

        if method == cv2.TM_SQDIFF_NORMED:
            pred = np.argmin(np.array(res))
        else:
            pred = np.argmax(np.array(res))

        all += 1
        if pred == gt:
            right += 1

        img_name = test_img_dir.split('/')[-1]
        print(f'for img {img_name}, pred is {pred}, gt is {gt}.')
    print(f'total sample number is {all}, and the accuracy is equal to: {right/all}')